{-# OPTIONS_GHC -Wall -Werror #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeSynonymInstances #-}

module StreamEngine 
( Occurrence
, Incident
, Signal
, Event
, Trigger
, (~<~)
, (~<=~)
, (~>~)
, (~>=~)
, (~&&~)
, (~||~)
, ite
, step
, edge
, edgeT
, trigger
, rising
, falling
, boolToMaybe
, maybeToBool
, boolToIncident
, integral
, bothO
, bothI
, bothE
, bothT
, eitherO
, eitherI
, eitherE
, eitherT
, mergeO
, mergeI
, mergeE
, mergeT
, mapS
, mapO
, mapI
, mapE
, mapT
, tag
, tag_
, section
, switch
, dump
, stamp
, traceI
, traceO
, traceE
, traceT
, traceS
)
where

import Math

import Data.Maybe
import Data.List

import Graphics.UI.GLFW (time)
import Graphics.Rendering.OpenGL (get)

import Control.Applicative
import Control.Monad

import Debug.Trace

import System.IO.Unsafe

import Data.VectorSpace hiding (Scalar)
import qualified Data.VectorSpace as VS

import qualified Data.Stream as S
import Data.Stream (Stream(..),(<:>))

-------------------------------------------------------------------------------

type Occurrence a = Maybe a
type Incident = Occurrence ()

-------------------------------------------------------------------------------

type Signal a = Stream a
type Event a  = Stream (Occurrence a)
type Trigger  = Stream Incident

-------------------------------------------------------------------------------
-- Arithmetic lifted to streams.

instance AdditiveGroup v => AdditiveGroup (Signal v) where
  zeroV = pure zeroV
  (^+^) = liftA2 (^+^)
  negateV = liftA negateV

instance VectorSpace v => VectorSpace (Signal v) where
  type VS.Scalar (Signal v) = Signal (VS.Scalar v)
  (*^) = liftA2 (*^)

instance Num n => Num (Signal n) where
  (+) = liftA2 (+)
  (*) = liftA2 (*)
  abs = liftA abs
  signum = liftA signum
  fromInteger = pure . fromInteger

instance Fractional n => Fractional (Signal n) where
  fromRational = pure . fromRational

instance Floating n => Floating (Signal n) where
  pi = pure pi
  exp = liftA exp
  sqrt = liftA sqrt
  log = liftA log
  (**) = liftA2 (**)
  logBase = liftA2 logBase
  sin = liftA sin
  tan = liftA tan
  cos = liftA cos
  asin = liftA asin
  atan = liftA atan
  acos = liftA acos
  sinh = liftA sinh
  tanh = liftA tanh
  cosh = liftA cosh
  asinh = liftA asinh
  atanh = liftA atanh
  acosh = liftA acosh

-------------------------------------------------------------------------------
-- Comparison and conditional lifted to streams.

infixr 4 ~<~ 
infixr 4 ~<=~
infixr 4 ~>~ 
infixr 4 ~>=~

(~&&~) :: Signal Bool -> Signal Bool -> Signal Bool
(~&&~) = liftA2 (&&)

(~||~) :: Signal Bool -> Signal Bool -> Signal Bool
(~||~) = liftA2 (||)

infixr 3 ~&&~
infixr 3 ~||~

(~<~) :: Ord a => Signal a -> Signal a -> Signal Bool
(~<~)  = liftA2 (<)

(~<=~) :: Ord a => Signal a -> Signal a -> Signal Bool
(~<=~) = liftA2 (<=)

(~>~) :: Ord a => Signal a -> Signal a -> Signal Bool
(~>~)  = liftA2 (>)

(~>=~) :: Ord a => Signal a -> Signal a -> Signal Bool
(~>=~) = liftA2 (>=)

ite :: Signal Bool -> Signal a -> Signal a -> Signal a
ite = liftA3 aux
  where
    aux b t f = if b then t else f

boolToMaybe :: a -> Bool -> Maybe a
boolToMaybe _ False = Nothing
boolToMaybe a True  = Just a

maybeToBool :: Maybe a -> Bool
maybeToBool Nothing = False
maybeToBool (Just _) = True

boolToIncident :: Bool -> Incident
boolToIncident = boolToMaybe ()

trigger :: Event a -> Trigger
trigger = (const () `mapE`)

-------------------------------------------------------------------------------

step :: a -> Event a -> Signal a
step firstValue eventStream = firstValue <:> rest eventStream
  where
    rest (maybeNextValue `Cons` nextEventStream) =
      step (maybe firstValue id maybeNextValue) nextEventStream

edge :: Eq a => a -> Signal a -> Event (a,a)
edge firstValue signal = S.zipWith aux (firstValue <:> signal) signal
  where
    aux a1 a2 = boolToMaybe (a1,a2) (a1 /= a2)

edgeT :: Eq a => a -> Signal a -> Trigger
edgeT firstValue = trigger . edge firstValue

rising :: Signal Bool -> Trigger
rising = fmap onlyRising . edge False
  where
    onlyRising (Just (False,True)) = Just ()
    onlyRising _                   = Nothing

falling :: Signal Bool -> Trigger
falling = rising . fmap not

-------------------------------------------------------------------------------
-- Transpose a list of signals into a signal of lists
-- This is like taking the cross section of the signals.

section :: [Signal a] -> Signal [a]
section = S.fromList . transpose . map S.toList

-------------------------------------------------------------------------------
-- Just an euler integral. Should add higher order ones.

integral :: VectorSpace a => a -> Signal a -> Signal (VS.Scalar a) -> Signal a
integral x vs dts = S.scan' (^+^) x (vs ^* dts)

-------------------------------------------------------------------------------
-- TODO: Not sure these seqs are needed.

bothO :: Occurrence a -> Occurrence b -> Occurrence (a,b)
bothO Nothing y         = y `seq` Nothing
bothO x Nothing         = x `seq` Nothing
bothO (Just x) (Just y) = Just (x,y)

bothI :: Incident -> Incident -> Incident
bothI Nothing y         = y `seq` Nothing
bothI x Nothing         = x `seq` Nothing
bothI (Just _) (Just _) = Just ()

eitherO :: Occurrence a -> Occurrence b -> Occurrence (Either a b)
eitherO (Just x) y = y `seq` Just $ Left x
eitherO x (Just y) = x `seq` Just $ Right y
eitherO x y        = x `seq` y `seq` Nothing

mergeO :: Occurrence a -> Occurrence a -> Occurrence a
mergeO (Just x) y = y `seq` Just x
mergeO x (Just y) = x `seq` Just y
mergeO x y        = x `seq` y`seq` Nothing

eitherI :: Incident -> Incident -> Incident
eitherI = mergeO

mergeI :: Incident -> Incident -> Incident
mergeI = mergeO

bothE :: Event a -> Event b -> Event (a,b)
bothE = liftA2 bothO

bothT :: Trigger -> Trigger -> Trigger
bothT = liftA2 bothI

eitherE :: Event a -> Event b -> Event (Either a b)
eitherE = liftA2 eitherO

mergeE :: Event a -> Event a -> Event a
mergeE = liftA2 mergeO

mergeT :: Trigger -> Trigger -> Trigger
mergeT = mergeE

eitherT :: Trigger -> Trigger -> Trigger
eitherT = mergeE

tag :: Event a -> Signal b -> Event (a,b)
tag (Just a  `Cons` as) (b `Cons` bs) = Just (a,b) `Cons` tag as bs
tag (Nothing `Cons` as) (_ `Cons` bs) = Nothing    `Cons` tag as bs

tag_ :: Event a -> Signal b -> Event b
tag_ as bs = fmap snd <$> tag as bs
 
-------------------------------------------------------------------------------
-- TODO: Refactor this, Signal carrier should not be needed.

{-
switch :: Signal a -> 
          Event (a -> Signal carrier -> Signal a) ->
          Signal carrier -> Signal a
switch (a `Cons` as) ~(mf `Cons` mfs) (_ `Cons` ss) =  
  a `Cons` maybe (switch as mfs ss) (\f -> f a ss) mf
-}

switch :: Signal a -> Event (a -> Signal a -> Signal a) -> Signal a
switch (a `Cons` as) ~(mf `Cons` mfs) =  
  a `Cons` maybe (switch as mfs) (\f -> f a as) mf

-------------------------------------------------------------------------------
-- Functors helpers

mapS :: (a->b) -> Signal a -> Signal b
mapS = fmap

mapO :: (a->b) -> Occurrence a -> Occurrence b
mapO = fmap

mapE :: (a->b) -> Event a -> Event b
mapE = fmap . fmap

mapI :: a -> Incident -> Occurrence a 
mapI = fmap . const

mapT :: a -> Trigger -> Event a
mapT = fmap . mapI

-------------------------------------------------------------------------------
-- Debugging

dump :: Show a => a -> a
dump x = trace (show x) x

stamp :: a -> (Time,a)
stamp x = unsafePerformIO $ do
            t <- get time
            return (t,x)

traceS :: String -> Signal a -> Signal a
traceS = fmap . trace . show . stamp

traceO :: String -> Occurrence a -> Occurrence a
traceO = fmap . trace . show . stamp

traceI :: String -> Incident -> Incident
traceI = traceO

traceE :: String -> Event a -> Event a
traceE = fmap . traceO

traceT :: String -> Trigger -> Trigger
traceT = traceE

-------------------------------------------------------------------------------
